from pyVim.connect import SmartConnectNoSSL, Disconnect
from pyVmomi import vim
import atexit
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv
import os

load_dotenv()  # 加载 .env 文件

VCENTER_HOST = os.getenv("VCENTER_HOST")
VCENTER_USER = os.getenv("VCENTER_USER")
VCENTER_PASSWORD = os.getenv("VCENTER_PASSWORD")

mcp = FastMCP(
    name="power_off_tool",
    description="安全关闭指定名称的主机",
    version="1.0.0",
)

def WaitForTask(task):
    while task.info.state not in [vim.TaskInfo.State.success, vim.TaskInfo.State.error]:
        pass
    if task.info.state == vim.TaskInfo.State.error:
        raise task.info.error
    
    return task.info.state

def power_off_host(host_name: str):
    """
    安全关闭指定名称的主机
    
    参数:
        host_name (str): 要关闭的vSphere主机名称
        
    返回:
        str: 操作结果消息
        
    流程:
        1. 连接到vCenter服务器
        2. 查找指定名称的主机
        3. 检查主机状态是否允许关闭
        4. 检查主机上是否有运行中的虚拟机
        5. 执行主机关闭操作
        6. 返回操作结果
    """
    si = SmartConnectNoSSL(host=VCENTER_HOST, user=VCENTER_USER, pwd=VCENTER_PASSWORD, port=443)
    atexit.register(Disconnect, si)

    content = si.RetrieveContent()
    container = content.viewManager.CreateContainerView(content.rootFolder, [vim.HostSystem], True)
    hosts = container.view

    host = None
    for host_instance in hosts:
        if host_instance.name == host_name:
            host = host_instance
            break

    if not host:
        return f"未能找到名称为 {host_name} 的主机"

    if host.summary.runtime.connectionState != vim.HostSystem.ConnectionState.connected:
        return f"主机 {host.name} 当前连接状态不允许关闭操作"

    vms = host.vm
    for vm in vms:
        if vm.summary.runtime.powerState == vim.VirtualMachinePowerState.poweredOn:
            return f"主机 {host.name} 上有虚拟机 {vm.name} 处于开启状态，无法安全关闭主机"

    if host.runtime.powerState != vim.HostSystem.PowerState.poweredOn:
        return f"主机 {host.name} 当前未开机无法执行操作"

    task = host.ShutdownHost_Task(force=True)
    result = WaitForTask(task)
    if result == vim.TaskInfo.State.success:
        return f"主机 {host.name} 已成功安全关闭"
    else:
        return f"关闭主机 {host.name} 时发生错误"

